#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
自动解析 air_to_sea_log.txt 并绘制所有主要物理量曲线
用法:python3 plot_air_to_sea_log.py air_to_sea_log.txt
"""
import re
import sys
import matplotlib.pyplot as plt
import numpy as np
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['Noto Sans CJK SC', 'SimHei', 'WenQuanYi Zen Hei']
matplotlib.rcParams['axes.unicode_minus'] = False

# 正则表达式
imu_re = re.compile(r"\[IMU\] t=([\d.]+)s 线加速度: ax=([\-\d.e]+), ay=([\-\d.e]+), az=([\-\d.e]+)\s+角速度: gx=([\-\d.e]+), gy=([\-\d.e]+), gz=([\-\d.e]+)")
odom_re = re.compile(r"\[ODOM\] t=([\d.]+)s 位置: x=([\-\d.e]+), y=([\-\d.e]+), z=([\-\d.e]+)\s+线速度: vx=([\-\d.e]+), vy=([\-\d.e]+), vz=([\-\d.e]+)\s+线加速度: ax=([\-\d.e]+), ay=([\-\d.e]+), az=([\-\d.e]+)\s+角速度: wx=([\-\d.e]+), wy=([\-\d.e]+), wz=([\-\d.e]+)\s+流体密度: ([\d.e]+) 阻尼缩放: ([\d.e]+)")

def parse_log(filename):
    imu_data = []
    odom_data = []
    with open(filename, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    i = 0
    while i < len(lines):
        line = lines[i].strip()
        if line.startswith('[IMU]'):
            # 合并下一行
            if i+1 < len(lines):
                line2 = lines[i+1].strip()
                full_line = line + ' ' + line2
                m = imu_re.match(full_line)
                if m:
                    imu_data.append([float(x) for x in m.groups()])
            i += 2
        elif line.startswith('[ODOM]'):
            # 合并后4行
            full_line = line
            for j in range(1,5):
                if i+j < len(lines):
                    full_line += ' ' + lines[i+j].strip()
            m = odom_re.match(full_line)
            if m:
                odom_data.append([float(x) for x in m.groups()])
            i += 5
        else:
            i += 1
    return np.array(imu_data), np.array(odom_data)

def plot_all(imu_data, odom_data):
    # IMU: t, ax, ay, az, gx, gy, gz
    # ODOM: t, x, y, z, vx, vy, vz, ax, ay, az, wx, wy, wz, density, damping
    plt.figure(figsize=(16, 10))
    # 1. IMU加速度
    plt.subplot(3,3,1)
    plt.title('IMU 线加速度')
    plt.plot(imu_data[:,0], imu_data[:,1], label='ax')
    plt.plot(imu_data[:,0], imu_data[:,2], label='ay')
    plt.plot(imu_data[:,0], imu_data[:,3], label='az')
    plt.xlabel('时间/s')
    plt.ylabel('加速度 (m/s²)')
    plt.legend()
    # 2. IMU角速度
    plt.subplot(3,3,2)
    plt.title('IMU 角速度')
    plt.plot(imu_data[:,0], imu_data[:,4], label='gx')
    plt.plot(imu_data[:,0], imu_data[:,5], label='gy')
    plt.plot(imu_data[:,0], imu_data[:,6], label='gz')
    plt.xlabel('时间/s')
    plt.ylabel('角速度 (rad/s)')
    plt.legend()
    # 3. ODOM位置
    plt.subplot(3,3,3)
    plt.title('ODOM 位置')
    plt.plot(odom_data[:,0], odom_data[:,1], label='x')
    plt.plot(odom_data[:,0], odom_data[:,2], label='y')
    plt.plot(odom_data[:,0], odom_data[:,3], label='z')
    plt.xlabel('时间/s')
    plt.ylabel('位置 (m)')
    plt.legend()
    # 4. ODOM线速度
    plt.subplot(3,3,4)
    plt.title('ODOM 线速度')
    plt.plot(odom_data[:,0], odom_data[:,4], label='vx')
    plt.plot(odom_data[:,0], odom_data[:,5], label='vy')
    plt.plot(odom_data[:,0], odom_data[:,6], label='vz')
    plt.xlabel('时间/s')
    plt.ylabel('速度 (m/s)')
    plt.legend()
    # 5. ODOM线加速度
    plt.subplot(3,3,5)
    plt.title('ODOM 线加速度')
    plt.plot(odom_data[:,0], odom_data[:,7], label='ax')
    plt.plot(odom_data[:,0], odom_data[:,8], label='ay')
    plt.plot(odom_data[:,0], odom_data[:,9], label='az')
    plt.xlabel('时间/s')
    plt.ylabel('加速度 (m/s²)')
    plt.legend()
    # 6. ODOM角速度
    plt.subplot(3,3,6)
    plt.title('ODOM 角速度')
    plt.plot(odom_data[:,0], odom_data[:,10], label='wx')
    plt.plot(odom_data[:,0], odom_data[:,11], label='wy')
    plt.plot(odom_data[:,0], odom_data[:,12], label='wz')
    plt.xlabel('时间/s')
    plt.ylabel('角速度 (rad/s)')
    plt.legend()
    # 7. 流体密度
    plt.subplot(3,3,7)
    plt.title('流体密度')
    plt.plot(odom_data[:,0], odom_data[:,13], label='density')
    plt.xlabel('时间/s')
    plt.ylabel('密度 (kg/m³)')
    plt.legend()
    # 8. 阻尼缩放
    plt.subplot(3,3,8)
    plt.title('阻尼缩放')
    plt.plot(odom_data[:,0], odom_data[:,14], label='damping')
    plt.xlabel('时间/s')
    plt.ylabel('缩放系数')
    plt.legend()
    plt.tight_layout()
    plt.savefig('air_to_sea_log_all.png', dpi=200)
    plt.show()

def main():
    if len(sys.argv) < 2:
        print('用法: python3 plot_air_to_sea_log.py air_to_sea_log.txt')
        sys.exit(1)
    imu_data, odom_data = parse_log(sys.argv[1])
    if imu_data.size == 0 or odom_data.size == 0:
        print('未能正确解析数据，请检查日志格式')
        sys.exit(2)
    plot_all(imu_data, odom_data)

if __name__ == '__main__':
    main() 